from tools.virtual_key_code import get_key_state
from bdtime import tt
from tools.utils import Key
from copy import deepcopy
from bdtime import Time
import os
import json
from tools.remove_temp_file import remove_path
from tools.mouse_state_recorder import DefaultArgs
from tools.beep import beep

kk = Key()


class KeyStateRecorder:
    def __init__(self, frames=DefaultArgs.frames, tt=None):
        self.frames = frames
        self.save_dir = os.path.join(*['..', 'tempdir', 'recorded_states', 'key'])

        self.recorded_states = []
        self.recorded_states.insert(0, {'recorder_type': self.__class__.__name__})

        self.targets = ['w', 's', 'a', 'd', 'f', 'e', 'q', '1', '2', '3', '4', 'ctrl', 'shift', 'space', 'esc', 'f1',
                        'alt']
        self.states = dict(zip(self.targets, [0] * len(self.targets)))
        self.is_recording = False

        self.tt = tt

        self.init_states()

    def init_states(self):
        self.states = dict(zip(self.targets, [0] * len(self.targets)))

    def record(self, until='alt + x'):
        tt = self.tt if self.tt else Time()

        self.is_recording = True
        i = 0
        DefaultArgs.flag__start_screen_record = True
        while True:
            if not self.is_recording or tt.stop(until, raise_error=False):
                break

            if DefaultArgs.flag__exit:
                break

            next_record_time = DefaultArgs.record_time_schedule[i]
            while tt.now() < next_record_time:
                if DefaultArgs.use_sleep_wait:
                    tt.sleep(0.001)

            self.states.update({'i': i})
            self.states.update({'time': next_record_time})
            # self.states.update({'time': tt.now()})

            for key in self.targets:
                state = get_key_state(key)
                if state:
                    self.states[key] = 1
                else:
                    self.states[key] = 0

            self.recorded_states.append(deepcopy(self.states))

            i += 1

        self.is_recording = False
        return self.recorded_states

    def play(self, recorded_states=None, speed=1):
        if not recorded_states:
            recorded_states = self.recorded_states

        self.init_states()

        if speed != 1:
            assert speed, 'speed不能为0!'
            for state in recorded_states:
                state['time'] = state['time'] / speed

        DefaultArgs.flag__start_screen_record = True

        tt = Time()
        print(f'~~~~~~ start key_state_recorder.play, and start_screen_record,  tt.now: {tt.now()}')

        for state in recorded_states:
            time = state.get('time')
            if time is None:
                continue

            now = tt.now()

            while now < time:
                tt.sleep(0.001)
                now = tt.now()

            if tt.stop('alt + x', raise_error=False):
                print(f'*** stop {self.__class__.__name__}.play by user!')
                break

            # print(f'--- tt.now: {tt.now()}, time: {time: .3f} --- delta_t: {tt.now() - time: .3f}')
            # if  tt.exceed(30):
            #     break
            # continue

            for k, v in state.items():
                if k == 'time':
                    continue

                if v:
                    self.press(k)
                else:
                    self.release(k)

        DefaultArgs.flag__stop_screen_record = True
        print(f'~~~~~~ end key_state_recorder.play, tt.now: {tt.now()}')

    def press(self, key):
        if get_key_state(key):
            return

        # if self.states[key]:
        #     return

        kk.press(key)
        self.states[key] = 1

    def release(self, key):
        if not get_key_state(key):
            return

        # if not self.states[key]:
        #     return

        kk.release(key)
        self.states[key] = 0

    def get_save_file_path(self, file_name):
        file_path = os.path.join(self.save_dir, file_name)
        return file_path

    def stop_record(self):
        pass

    def load_record(self, file_name):
        file_path = os.path.join(self.save_dir, file_name)
        with open(file_path, 'r') as f:
            self.recorded_states = json.load(f)
        e = self.recorded_states[0]
        assert e['recorder_type'] == self.__class__.__name__

    def save_record(self, file_name: str = ""):
        os.makedirs(self.save_dir, exist_ok=True)
        data = json.dumps(self.recorded_states, indent=4)
        timestamp = tt.get_current_beijing_time_str(tt.common_date_time_formats.ms_int[-8:])

        file_name = file_name if file_name else f'record__{timestamp}.json'

        file_path = os.path.join(self.save_dir, file_name)
        with open(file_path, 'w+') as f:
            f.write(data)

        print(f'{self.__class__.__name__} --- save_record to --- {file_path}')
        return file_path

    def remove_all_recorded_states(self):
        remove_path(self.save_dir, keep_external_folder=True)


def run(save_file_name: (None, str) = "", play=False, remove_else_recoded=False, key_state_recorder=None, tt=None,
        play_target=False, end_call=None):
    if not key_state_recorder:
        tt = tt if tt else Time()
        key_state_recorder = KeyStateRecorder(tt=tt)
    else:
        tt = key_state_recorder.tt

    try:
        if play_target is True:
            play_target = save_file_name

        if play_target:
            key_state_recorder.load_record(play_target)

        else:
            key_state_recorder.record()

            if remove_else_recoded:
                key_state_recorder.remove_all_recorded_states()

            if save_file_name is not None:
                key_state_recorder.save_record(file_name=save_file_name)

        if play:
            tt.sleep(1)
            print(f'====== start key_state_recorder.play, tt.now: {tt.now()}')
            key_state_recorder.play(speed=1)

        if end_call:
            end_call()

    except Exception as e:
        beep(400, 0.3)
        beep(600, 0.3)
        beep(800, 0.3)
        raise e


def main():
    import keyboard
    from tools.beep import beep

    # hotkey = 'alt + `'
    # print(f'--- wait hotkey[{hotkey}]!')
    # keyboard.wait(hotkey)

    tt.sleep(0.5)
    print('--- start main__mouse_recorder!')
    beep()

    # save_file_name = '蒙德传送点'
    save_file_name = 'test.json'
    play_target = True
    key_kwargs = {
        'save_file_name': save_file_name,
        'play': True,
        'remove_else_recoded': False,
        'play_target': play_target,
        'tt': tt
    }
    run(**key_kwargs)


if __name__ == '__main__':
    main()
